{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip -q install -U langgraph langchain langchain-google-genai google-generativeai typing_extensions\n",
        "!pip -q install \"requests==2.32.4\"\n",
        "\n",
        "import os\n",
        "import json\n",
        "import textwrap\n",
        "import getpass\n",
        "import time\n",
        "from typing import Annotated, List, Dict, Any, Optional\n",
        "\n",
        "from typing_extensions import TypedDict\n",
        "\n",
        "from langchain.chat_models import init_chat_model\n",
        "from langchain_core.messages import BaseMessage\n",
        "from langchain_core.tools import tool\n",
        "\n",
        "from langgraph.graph import StateGraph, START, END\n",
        "from langgraph.graph.message import add_messages\n",
        "from langgraph.checkpoint.memory import InMemorySaver\n",
        "from langgraph.prebuilt import ToolNode, tools_condition\n",
        "\n",
        "import requests\n",
        "from requests.adapters import HTTPAdapter, Retry\n",
        "\n",
        "if not os.environ.get(\"GOOGLE_API_KEY\"):\n",
        "    os.environ[\"GOOGLE_API_KEY\"] = getpass.getpass(\"🔑 Enter your Google API Key (Gemini): \")\n",
        "\n",
        "llm = init_chat_model(\"google_genai:gemini-2.0-flash\")"
      ],
      "metadata": {
        "id": "M5DVj4WSMrZ3"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "WIKI_SEARCH_URL = \"https://en.wikipedia.org/w/api.php\"\n",
        "\n",
        "_session = requests.Session()\n",
        "_session.headers.update({\n",
        "    \"User-Agent\": \"LangGraph-Colab-Demo/1.0 (contact: example@example.com)\",\n",
        "    \"Accept\": \"application/json\",\n",
        "})\n",
        "retry = Retry(\n",
        "    total=5, connect=5, read=5, backoff_factor=0.5,\n",
        "    status_forcelist=(429, 500, 502, 503, 504),\n",
        "    allowed_methods=(\"GET\", \"POST\")\n",
        ")\n",
        "_session.mount(\"https://\", HTTPAdapter(max_retries=retry))\n",
        "_session.mount(\"http://\", HTTPAdapter(max_retries=retry))\n",
        "\n",
        "def _wiki_search_raw(query: str, limit: int = 3) -> List[Dict[str, str]]:\n",
        "    \"\"\"\n",
        "    Use MediaWiki search API with:\n",
        "      - origin='*' (good practice for CORS)\n",
        "      - Polite UA + retries\n",
        "    Returns compact list of {title, snippet_html, url}.\n",
        "    \"\"\"\n",
        "    params = {\n",
        "        \"action\": \"query\",\n",
        "        \"list\": \"search\",\n",
        "        \"format\": \"json\",\n",
        "        \"srsearch\": query,\n",
        "        \"srlimit\": limit,\n",
        "        \"srprop\": \"snippet\",\n",
        "        \"utf8\": 1,\n",
        "        \"origin\": \"*\",\n",
        "    }\n",
        "    r = _session.get(WIKI_SEARCH_URL, params=params, timeout=15)\n",
        "    r.raise_for_status()\n",
        "    data = r.json()\n",
        "    out = []\n",
        "    for item in data.get(\"query\", {}).get(\"search\", []):\n",
        "        title = item.get(\"title\", \"\")\n",
        "        page_url = f\"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}\"\n",
        "        snippet = item.get(\"snippet\", \"\")\n",
        "        out.append({\"title\": title, \"snippet_html\": snippet, \"url\": page_url})\n",
        "    return out\n",
        "\n",
        "@tool\n",
        "def wiki_search(query: str) -> List[Dict[str, str]]:\n",
        "    \"\"\"Search Wikipedia and return up to 3 results with title, snippet_html, and url.\"\"\"\n",
        "    try:\n",
        "        results = _wiki_search_raw(query, limit=3)\n",
        "        return results if results else [{\"title\": \"No results\", \"snippet_html\": \"\", \"url\": \"\"}]\n",
        "    except Exception as e:\n",
        "        return [{\"title\": \"Error\", \"snippet_html\": str(e), \"url\": \"\"}]\n",
        "\n",
        "TOOLS = [wiki_search]"
      ],
      "metadata": {
        "id": "UCwMw0BGMyNI"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class State(TypedDict):\n",
        "    messages: Annotated[list, add_messages]\n",
        "\n",
        "graph_builder = StateGraph(State)\n",
        "\n",
        "llm_with_tools = llm.bind_tools(TOOLS)\n",
        "\n",
        "SYSTEM_INSTRUCTIONS = textwrap.dedent(\"\"\"\n",
        "You are ResearchBuddy, a careful research assistant.\n",
        "- If the user asks you to \"research\", \"find info\", \"latest\", \"web\", or references a library/framework/product,\n",
        "  you SHOULD call the `wiki_search` tool at least once before finalizing your answer.\n",
        "- When you call tools, be concise in the text you produce around the call.\n",
        "- After receiving tool results, cite at least the page titles you used in your summary.\n",
        "\"\"\").strip()\n",
        "\n",
        "def chatbot(state: State) -> Dict[str, Any]:\n",
        "    \"\"\"Single step: call the LLM (with tools bound) on the current messages.\"\"\"\n",
        "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
        "\n",
        "graph_builder.add_node(\"chatbot\", chatbot)\n",
        "\n",
        "tool_node = ToolNode(tools=TOOLS)\n",
        "graph_builder.add_node(\"tools\", tool_node)\n",
        "\n",
        "graph_builder.add_conditional_edges(\"chatbot\", tools_condition)\n",
        "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
        "graph_builder.add_edge(START, \"chatbot\")\n",
        "\n",
        "memory = InMemorySaver()\n",
        "graph = graph_builder.compile(checkpointer=memory)"
      ],
      "metadata": {
        "id": "yVGmCX2zM4Hh"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def print_last_message(event: Dict[str, Any]):\n",
        "    \"\"\"Pretty-print the last message in an event if available.\"\"\"\n",
        "    if \"messages\" in event and event[\"messages\"]:\n",
        "        msg = event[\"messages\"][-1]\n",
        "        try:\n",
        "            if isinstance(msg, BaseMessage):\n",
        "                msg.pretty_print()\n",
        "            else:\n",
        "                role = msg.get(\"role\", \"unknown\")\n",
        "                content = msg.get(\"content\", \"\")\n",
        "                print(f\"\\n[{role.upper()}]\\n{content}\\n\")\n",
        "        except Exception:\n",
        "            print(str(msg))\n",
        "\n",
        "def show_state_history(cfg: Dict[str, Any]) -> List[Any]:\n",
        "    \"\"\"Print a concise view of checkpoints; return the list as well.\"\"\"\n",
        "    history = list(graph.get_state_history(cfg))\n",
        "    print(\"\\n=== 📜 State history (most recent first) ===\")\n",
        "    for i, st in enumerate(history):\n",
        "        n = st.next\n",
        "        n_txt = f\"{n}\" if n else \"()\"\n",
        "        print(f\"{i:02d}) NumMessages={len(st.values.get('messages', []))}  Next={n_txt}\")\n",
        "    print(\"=== End history ===\\n\")\n",
        "    return history\n",
        "\n",
        "def pick_checkpoint_by_next(history: List[Any], node_name: str = \"tools\") -> Optional[Any]:\n",
        "    \"\"\"Pick the first checkpoint whose `next` includes a given node (e.g., 'tools').\"\"\"\n",
        "    for st in history:\n",
        "        nxt = tuple(st.next) if st.next else tuple()\n",
        "        if node_name in nxt:\n",
        "            return st\n",
        "    return None"
      ],
      "metadata": {
        "id": "qUjkLK1aM_X3"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = {\"configurable\": {\"thread_id\": \"demo-thread-1\"}}\n",
        "\n",
        "first_turn = {\n",
        "    \"messages\": [\n",
        "        {\"role\": \"system\", \"content\": SYSTEM_INSTRUCTIONS},\n",
        "        {\"role\": \"user\", \"content\": \"I'm learning LangGraph. Could you do some research on it for me?\"},\n",
        "    ]\n",
        "}\n",
        "\n",
        "print(\"\\n==================== 🟢 STEP 1: First user turn ====================\")\n",
        "events = graph.stream(first_turn, config, stream_mode=\"values\")\n",
        "for ev in events:\n",
        "    print_last_message(ev)\n",
        "\n",
        "second_turn = {\n",
        "    \"messages\": [\n",
        "        {\"role\": \"user\", \"content\": \"Ya that's helpful. Maybe I'll build an autonomous agent with it!\"}\n",
        "    ]\n",
        "}\n",
        "\n",
        "print(\"\\n==================== 🟢 STEP 2: Second user turn ====================\")\n",
        "events = graph.stream(second_turn, config, stream_mode=\"values\")\n",
        "for ev in events:\n",
        "    print_last_message(ev)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QXl-DwKDNCNV",
        "outputId": "87426463-6800-4875-8c14-138e6cec9ffe"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "==================== 🟢 STEP 1: First user turn ====================\n",
            "================================\u001b[1m Human Message \u001b[0m=================================\n",
            "\n",
            "I'm learning LangGraph. Could you do some research on it for me?\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "Tool Calls:\n",
            "  wiki_search (7c30a0a2-276f-4559-b3cd-718850d129b4)\n",
            " Call ID: 7c30a0a2-276f-4559-b3cd-718850d129b4\n",
            "  Args:\n",
            "    query: LangGraph\n",
            "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
            "Name: wiki_search\n",
            "\n",
            "[{\"title\": \"LangChain\", \"snippet_html\": \"launched <span class=\\\"searchmatch\\\">LangGraph</span> Platform into general availability, providing managed infrastructure for deploying long-running, stateful AI agents. <span class=\\\"searchmatch\\\">Lang</span>Chain&#039;s developers\", \"url\": \"https://en.wikipedia.org/wiki/LangChain\"}, {\"title\": \"C (programming language)\", \"snippet_html\": \"(PDF). Archived (PDF) from the original on October 25, 2007. (3.61 MB) comp.<span class=\\\"searchmatch\\\">lang</span>.c Frequently Asked Questions A History of C, by Dennis Ritchie C Library\", \"url\": \"https://en.wikipedia.org/wiki/C_(programming_language)\"}, {\"title\": \"Pharmacokinetics of estradiol\", \"snippet_html\": \"dienanthate, and undecylate, as well as polyestradiol phosphate—for more <span class=\\\"searchmatch\\\">graphs</span>. Estradiol esters like estradiol valerate and estradiol cypionate can be\", \"url\": \"https://en.wikipedia.org/wiki/Pharmacokinetics_of_estradiol\"}]\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "LangGraph is a platform launched into general availability that provides managed infrastructure for deploying long-running, stateful AI agents (LangChain). According to the search results, LangGraph is related to LangChain (\"LangChain\" Wikipedia page).\n",
            "\n",
            "==================== 🟢 STEP 2: Second user turn ====================\n",
            "================================\u001b[1m Human Message \u001b[0m=================================\n",
            "\n",
            "Ya that's helpful. Maybe I'll build an autonomous agent with it!\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "That sounds like a cool project! Is there anything else I can help you with regarding LangGraph or autonomous agents?\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "im9dfJkvK9ab",
        "outputId": "a0165854-b94d-429e-8c7e-dd00005118c6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "==================== 🔁 REPLAY: Full state history ====================\n",
            "\n",
            "=== 📜 State history (most recent first) ===\n",
            "00) NumMessages=7  Next=()\n",
            "01) NumMessages=6  Next=('chatbot',)\n",
            "02) NumMessages=5  Next=('__start__',)\n",
            "03) NumMessages=5  Next=()\n",
            "04) NumMessages=4  Next=('chatbot',)\n",
            "05) NumMessages=3  Next=('tools',)\n",
            "06) NumMessages=2  Next=('chatbot',)\n",
            "07) NumMessages=0  Next=('__start__',)\n",
            "=== End history ===\n",
            "\n",
            "Chosen checkpoint to resume from:\n",
            "  Next: ('tools',)\n",
            "  Config: {'configurable': {'thread_id': 'demo-thread-1', 'checkpoint_ns': '', 'checkpoint_id': '1f086431-27ee-6bd1-8001-f0def998f65f'}}\n",
            "\n",
            "==================== ⏪ RESUME from chosen checkpoint ====================\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "Tool Calls:\n",
            "  wiki_search (7c30a0a2-276f-4559-b3cd-718850d129b4)\n",
            " Call ID: 7c30a0a2-276f-4559-b3cd-718850d129b4\n",
            "  Args:\n",
            "    query: LangGraph\n",
            "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
            "Name: wiki_search\n",
            "\n",
            "[{\"title\": \"LangChain\", \"snippet_html\": \"launched <span class=\\\"searchmatch\\\">LangGraph</span> Platform into general availability, providing managed infrastructure for deploying long-running, stateful AI agents. <span class=\\\"searchmatch\\\">Lang</span>Chain&#039;s developers\", \"url\": \"https://en.wikipedia.org/wiki/LangChain\"}, {\"title\": \"C (programming language)\", \"snippet_html\": \"(PDF). Archived (PDF) from the original on October 25, 2007. (3.61 MB) comp.<span class=\\\"searchmatch\\\">lang</span>.c Frequently Asked Questions A History of C, by Dennis Ritchie C Library\", \"url\": \"https://en.wikipedia.org/wiki/C_(programming_language)\"}, {\"title\": \"Pharmacokinetics of estradiol\", \"snippet_html\": \"dienanthate, and undecylate, as well as polyestradiol phosphate—for more <span class=\\\"searchmatch\\\">graphs</span>. Estradiol esters like estradiol valerate and estradiol cypionate can be\", \"url\": \"https://en.wikipedia.org/wiki/Pharmacokinetics_of_estradiol\"}]\n",
            "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
            "\n",
            "LangGraph is a platform for deploying long-running, stateful AI agents. It provides managed infrastructure and was launched into general availability by the developers of LangChain (\"LangChain\").\n",
            "\n",
            "✅ Done. You added steps, replayed history, and resumed from a prior checkpoint.\n"
          ]
        }
      ],
      "source": [
        "print(\"\\n==================== 🔁 REPLAY: Full state history ====================\")\n",
        "history = show_state_history(config)\n",
        "\n",
        "to_replay = pick_checkpoint_by_next(history, node_name=\"tools\")\n",
        "if to_replay is None:\n",
        "    to_replay = history[min(2, len(history) - 1)]\n",
        "\n",
        "print(\"Chosen checkpoint to resume from:\")\n",
        "print(\"  Next:\", to_replay.next)\n",
        "print(\"  Config:\", to_replay.config)\n",
        "\n",
        "print(\"\\n==================== ⏪ RESUME from chosen checkpoint ====================\")\n",
        "for ev in graph.stream(None, to_replay.config, stream_mode=\"values\"):\n",
        "    print_last_message(ev)\n",
        "\n",
        "MANUAL_INDEX = None\n",
        "if MANUAL_INDEX is not None and 0 <= MANUAL_INDEX < len(history):\n",
        "    chosen = history[MANUAL_INDEX]\n",
        "    print(f\"\\n==================== 🧭 MANUAL RESUME @ index {MANUAL_INDEX} ====================\")\n",
        "    print(\"Next:\", chosen.next)\n",
        "    print(\"Config:\", chosen.config)\n",
        "    for ev in graph.stream(None, chosen.config, stream_mode=\"values\"):\n",
        "        print_last_message(ev)\n",
        "\n",
        "print(\"\\n✅ Done. You added steps, replayed history, and resumed from a prior checkpoint.\")"
      ]
    }
  ]
}